"""
基于IC趋势模型的股指策略
使用日线K线，趋势跟踪策略
"""

import numpy as np
import pandas as pd
from datetime import datetime
from tqsdk import TqApi, TargetPosTask

class ICTrendStrategy:
    strategy_name = "ICTrendStrategy"
    """
    IC趋势策略
    
    基于开盘到收盘的扩展价格参数指标，确定多空信号:
    1. 计算RET指标: HIGH + LOW - OPEN - CLOSE
    2. 分析近期RET的累计值和正负收益比
    3. 根据趋势强度生成多空信号
    
    策略逻辑:
    - 趋势向上且强度足够: 做多
    - 趋势向下且强度足够: 做空
    - 趋势不明确: 空仓
    """
    
    def __init__(self, api, symbol="IF", exchange="CFFEX", **kwargs):
        self.api = api
        self.symbol = symbol
        self.i_symbol = f"KQ.i@{exchange}.{symbol}"
        self.m_symbol = f"KQ.m@{exchange}.{symbol}"
        
        # IC趋势策略参数
        self.alpha1_var = kwargs.get('alpha1_var', 3)  # RET累计周期
        self.lookback_period = kwargs.get('lookback_period', 6)  # 回看周期
        self.start_period = kwargs.get('start_period', 20)  # 策略启动需要的最小数据量
        
        # 风控参数
        self.market_period = kwargs.get('market_period', 60)
        self.trading_close = kwargs.get('trading_close', False)  # 交易K线的位置
        
        # 数据订阅 - 使用日线K线
        self.kline_daily = api.get_kline_serial(self.i_symbol, self.market_period * 60, data_length=800)
        self.kline = None
        
        # 交易对象
        self.main_symbol = api.get_quote(self.m_symbol).underlying_symbol
        self.target_pos = TargetPosTask(api, self.main_symbol)
        symbol_info = api.get_quote(self.main_symbol)
        symbol_value = symbol_info.volume_multiple * symbol_info.last_price
        symbol_margin = symbol_info.margin
        self.position_size = int(api.get_account().balance/symbol_value)
        
        # 交易状态
        self.last_signal = 0  # 上一次信号: 1=多头, -1=空头, 0=空仓
        self.last_trade_time = 0
        self.ret_series = []  # RET指标序列
        
    def calculate_ret_indicator(self) -> float:
        """
        计算RET趋势指标
        
        Returns:
            当前的信号值: 1=做多, -1=做空, 0=空仓
        """
        try:
            if len(self.kline_daily) < self.start_period:
                return 0

            
            # 获取OHLC数据
            open_prices = self.klines.open.values
            high_prices = self.klines.high.values
            low_prices = self.klines.low.values
            close_prices = self.klines.close.values
            
            # 计算RET指标
            current_idx = len(open_prices) - 1

            # 更新RET序列
            if len(self.ret_series) >= 1:  # 保持序列长度
                ret_value = (high_prices[-1] + low_prices[-1] - open_prices[-1] - close_prices[-1])
                self.ret_series[:-1] = self.ret_series[1:]
                self.ret_series[-1] = ret_value
            else:
                self.ret_series =high_prices + low_prices - open_prices - close_prices
            
            # 获取分析窗口
            ANSX = self.ret_series[-self.lookback_period:]
            
            # 计算正负收益的均值
            positive_mean = np.mean(ANSX[ANSX > 0]) if len(ANSX[ANSX > 0]) > 0 else 0
            negative_mean = np.mean(ANSX[ANSX < 0]) if len(ANSX[ANSX < 0]) > 0 else 0
            
            # 计算近期RET累计值
            sum_ret = sum(self.ret_series[-(self.alpha1_var + 1):])

            # 生成信号
            signal = 0
            
            if (sum_ret >= 0 and negative_mean != 0 and 
                abs(positive_mean) / abs(negative_mean) >= 1):
                signal = 1  # 做多信号
            elif (sum_ret < 0 and negative_mean != 0 and 
                  abs(positive_mean) / abs(negative_mean) < 1):
                signal = -1  # 做空信号
            else:
                 signal = 0  # 空仓信号
            
            return signal
            
        except Exception as e:
            print(f"计算RET指标失败: {e}")
            return 0
    
    def calculate_atr(self, period: int = 14) -> float:
        """计算ATR用于风控"""
        try:
            if len(self.klines) < period + 1:
                return self.klines.close.iloc[-1] * 0.02  # 默认2%
            
            high = self.klines.high.iloc[-period:]
            low = self.klines.low.iloc[-period:]
            close = self.klines.close.iloc[-period-1:-1]  # 前一日收盘价
            
            tr1 = high - low
            tr2 = abs(high - close.values)
            tr3 = abs(low - close.values)
            
            tr = np.maximum(tr1, np.maximum(tr2, tr3))
            atr = np.mean(tr)
            
            return atr
            
        except Exception as e:
            return klines.close.iloc[-1] * 0.02
    
    def on_bar(self):
        """主循环"""
        bar_count = 0
        while True:
            self.api.wait_update()
            if self.api.is_changing(self.kline_daily):
                bar_count += 1
                current_time = self.kline_daily.datetime.iloc[-1]
                current_price = self.kline_daily.close.iloc[-1]
                volume = self.kline_daily.volume.iloc[-1]
                
                # 每10根K线打印一次进度
                if bar_count % 10 == 0:
                    print(f"处理第{bar_count}根日K线: 时间={current_time}, 价格={current_price:.1f}")
                self.klines = self.kline_daily if self.trading_close else self.kline_daily.iloc[:-1]
                if self.trading_close:
                    if volume != 0:
                        self.run()
                else:
                    if volume == 0:
                        self.run()
    
    def run(self):
        """策略执行逻辑"""
        try:
            # 切换主力合约
            self.main_symbol = self.api.get_quote(self.m_symbol).underlying_symbol
            self.target_pos = TargetPosTask(self.api, self.main_symbol)
            
            # 平掉非当前主力的持仓
            positions = self.api.get_position()
            for sym, pos in positions.items():
                if sym != self.main_symbol and pos.pos != 0:
                    TargetPosTask(self.api, sym).set_target_volume(0)
            
            main_pos = positions.get(self.main_symbol, None)
            current_pos = main_pos.pos if main_pos else 0
            current_price = self.klines.close.iloc[-1]
            current_time = self.kline_daily.datetime.iloc[-1]
            
            # 检查数据是否足够
            if len(self.klines) < self.start_period:
                return
            
            # 控制交易频率 (日线策略，每日最多交易一次)
            # if current_time - self.last_trade_time < 24 * 60 * 60 * 1000:
            #     return
            
            # 计算信号
            current_signal = self.calculate_ret_indicator()
            
            # 调试信息：每10次运行打印一次
            if hasattr(self, 'run_count'):
                self.run_count += 1
            else:
                self.run_count = 1
            
            if self.run_count < self.start_period:
                return

            if self.run_count % 10 == 0:
                ret_value = self.ret_series[-1]
                print(f"{self.strategy_name}策略运行{self.run_count}次: RET={ret_value:.3f}, 信号={current_signal}, 持仓={current_pos}, 价格={current_price:.1f}")
            
            # 判断是否需要交易
            need_trade = False
            trade_reason = ""
            
            if current_pos == 0:  # 无持仓
                if current_signal == 1:  # 做多信号
                    self.target_pos.set_target_volume(self.position_size)
                    need_trade = True
                    trade_reason = "🟢 开多"
                    
                elif current_signal == -1:  # 做空信号
                    self.target_pos.set_target_volume(-self.position_size)
                    need_trade = True
                    trade_reason = "🔴 开空"
            
            else:  # 有持仓
                # 检查信号变化
                if current_signal != self.last_signal:
                    if current_signal == 1 and current_pos <= 0:  # 转为做多
                        self.target_pos.set_target_volume(self.position_size)
                        need_trade = True
                        trade_reason = "🟢 转多"
                        
                    elif current_signal == -1 and current_pos >= 0:  # 转为做空
                        self.target_pos.set_target_volume(-self.position_size)
                        need_trade = True
                        trade_reason = "🔴 转空"
                        
                    elif current_signal == 0:  # 平仓
                        self.target_pos.set_target_volume(0)
                        need_trade = True
                        trade_reason = "⚪ 平仓"
            
            # 执行交易
            if need_trade:
                self.last_trade_time = current_time
                self.last_signal = current_signal
                
                ret_value = self.ret_series[-1]
                print(f"{trade_reason}: RET={ret_value:.3f}, 信号={current_signal}, 价格={current_price:.1f}")
                
        except Exception as e:
            print(f"策略执行失败: {e}")

